#!/usr/bin/env python3
"""
性能测试脚本
测试系统的帧率、响应时间和资源使用情况
"""

import sys
import os
from pathlib import Path
import time
import psutil
import threading
import statistics

# 添加src目录到Python路径
current_dir = Path(__file__).parent
src_dir = current_dir / "src"
sys.path.insert(0, str(src_dir))

import cv2
import numpy as np
import logging

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class PerformanceMonitor:
    """性能监控器"""
    
    def __init__(self):
        self.frame_times = []
        self.detection_times = []
        self.cpu_usage = []
        self.memory_usage = []
        self.is_monitoring = False
        self.monitor_thread = None
    
    def start_monitoring(self):
        """开始监控"""
        self.is_monitoring = True
        self.monitor_thread = threading.Thread(target=self._monitor_resources)
        self.monitor_thread.daemon = True
        self.monitor_thread.start()
    
    def stop_monitoring(self):
        """停止监控"""
        self.is_monitoring = False
        if self.monitor_thread:
            self.monitor_thread.join()
    
    def _monitor_resources(self):
        """监控系统资源"""
        while self.is_monitoring:
            # CPU使用率
            cpu_percent = psutil.cpu_percent(interval=0.1)
            self.cpu_usage.append(cpu_percent)
            
            # 内存使用率
            memory = psutil.virtual_memory()
            self.memory_usage.append(memory.percent)
            
            time.sleep(0.5)
    
    def record_frame_time(self, frame_time):
        """记录帧处理时间"""
        self.frame_times.append(frame_time)
    
    def record_detection_time(self, detection_time):
        """记录检测时间"""
        self.detection_times.append(detection_time)
    
    def get_statistics(self):
        """获取统计信息"""
        stats = {}
        
        if self.frame_times:
            stats['frame_rate'] = {
                'avg_fps': 1.0 / statistics.mean(self.frame_times) if self.frame_times else 0,
                'min_fps': 1.0 / max(self.frame_times) if self.frame_times else 0,
                'max_fps': 1.0 / min(self.frame_times) if self.frame_times else 0,
                'avg_frame_time': statistics.mean(self.frame_times) * 1000,  # ms
            }
        
        if self.detection_times:
            stats['detection'] = {
                'avg_time': statistics.mean(self.detection_times) * 1000,  # ms
                'min_time': min(self.detection_times) * 1000,
                'max_time': max(self.detection_times) * 1000,
            }
        
        if self.cpu_usage:
            stats['cpu'] = {
                'avg_usage': statistics.mean(self.cpu_usage),
                'max_usage': max(self.cpu_usage),
            }
        
        if self.memory_usage:
            stats['memory'] = {
                'avg_usage': statistics.mean(self.memory_usage),
                'max_usage': max(self.memory_usage),
            }
        
        return stats

def test_camera_performance():
    """测试摄像头性能"""
    print("=" * 50)
    print("📹 摄像头性能测试")
    print("=" * 50)
    
    try:
        from camera_manager import CameraManager
        
        camera = CameraManager(camera_id=0, width=640, height=480)
        monitor = PerformanceMonitor()
        
        if not camera.initialize():
            print("❌ 摄像头初始化失败")
            return False
        
        print("✅ 摄像头初始化成功，开始性能测试...")
        
        monitor.start_monitoring()
        
        # 测试100帧
        test_frames = 100
        successful_frames = 0
        
        for i in range(test_frames):
            start_time = time.time()
            
            ret, frame = camera.cap.read()
            if ret:
                successful_frames += 1
                # 模拟一些处理
                gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                blurred = cv2.GaussianBlur(gray, (5, 5), 0)
            
            frame_time = time.time() - start_time
            monitor.record_frame_time(frame_time)
            
            if (i + 1) % 20 == 0:
                print(f"进度: {i + 1}/{test_frames} 帧")
        
        monitor.stop_monitoring()
        camera.release()
        
        # 获取统计信息
        stats = monitor.get_statistics()
        
        print(f"\n📊 摄像头性能统计:")
        print(f"成功帧数: {successful_frames}/{test_frames}")
        print(f"成功率: {successful_frames/test_frames*100:.1f}%")
        
        if 'frame_rate' in stats:
            fr = stats['frame_rate']
            print(f"平均帧率: {fr['avg_fps']:.1f} FPS")
            print(f"帧率范围: {fr['min_fps']:.1f} - {fr['max_fps']:.1f} FPS")
            print(f"平均帧时间: {fr['avg_frame_time']:.1f} ms")
        
        if 'cpu' in stats:
            print(f"CPU使用率: 平均 {stats['cpu']['avg_usage']:.1f}%, 峰值 {stats['cpu']['max_usage']:.1f}%")
        
        if 'memory' in stats:
            print(f"内存使用率: 平均 {stats['memory']['avg_usage']:.1f}%, 峰值 {stats['memory']['max_usage']:.1f}%")
        
        return successful_frames > test_frames * 0.8  # 80%成功率
        
    except Exception as e:
        print(f"❌ 摄像头性能测试失败: {e}")
        return False

def test_yolo_performance():
    """测试YOLO检测性能"""
    print("\n" + "=" * 50)
    print("🎯 YOLO检测性能测试")
    print("=" * 50)
    
    try:
        from yolo_detector import YOLODetector
        
        detector = YOLODetector(model_path="yolo11n.pt", confidence=0.5)
        monitor = PerformanceMonitor()
        
        if not detector.initialize():
            print("❌ YOLO模型初始化失败")
            return False
        
        print("✅ YOLO模型初始化成功，开始性能测试...")
        
        monitor.start_monitoring()
        
        # 创建测试图像
        test_images = []
        for i in range(50):
            # 创建不同的测试图像
            img = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
            # 添加一些形状
            cv2.rectangle(img, (100+i, 100+i), (200+i, 200+i), (255, 255, 255), -1)
            cv2.circle(img, (400-i, 300+i), 30, (0, 255, 0), -1)
            test_images.append(img)
        
        successful_detections = 0
        
        for i, img in enumerate(test_images):
            start_time = time.time()
            
            detections = detector.detect(img)
            successful_detections += 1
            
            detection_time = time.time() - start_time
            monitor.record_detection_time(detection_time)
            
            if (i + 1) % 10 == 0:
                print(f"进度: {i + 1}/{len(test_images)} 张图像")
        
        monitor.stop_monitoring()
        
        # 获取统计信息
        stats = monitor.get_statistics()
        
        print(f"\n📊 YOLO检测性能统计:")
        print(f"成功检测: {successful_detections}/{len(test_images)}")
        
        if 'detection' in stats:
            det = stats['detection']
            print(f"平均检测时间: {det['avg_time']:.1f} ms")
            print(f"检测时间范围: {det['min_time']:.1f} - {det['max_time']:.1f} ms")
            print(f"理论最大FPS: {1000/det['avg_time']:.1f}")
        
        if 'cpu' in stats:
            print(f"CPU使用率: 平均 {stats['cpu']['avg_usage']:.1f}%, 峰值 {stats['cpu']['max_usage']:.1f}%")
        
        if 'memory' in stats:
            print(f"内存使用率: 平均 {stats['memory']['avg_usage']:.1f}%, 峰值 {stats['memory']['max_usage']:.1f}%")
        
        return successful_detections == len(test_images)
        
    except Exception as e:
        print(f"❌ YOLO性能测试失败: {e}")
        return False

def test_integrated_performance():
    """集成性能测试"""
    print("\n" + "=" * 50)
    print("🔧 集成性能测试")
    print("=" * 50)
    
    try:
        from camera_manager import CameraManager
        from yolo_detector import YOLODetector
        from crosshair import Crosshair
        
        # 初始化组件
        camera = CameraManager(camera_id=0, width=640, height=480)
        detector = YOLODetector(model_path="yolo11n.pt", confidence=0.5)
        crosshair = Crosshair(frame_width=640, frame_height=480)
        monitor = PerformanceMonitor()
        
        # 检查摄像头是否可用
        if not camera.initialize():
            print("⚠️  摄像头不可用，使用模拟数据进行测试")
            use_camera = False
        else:
            use_camera = True
        
        if not detector.initialize():
            print("❌ YOLO模型初始化失败")
            return False
        
        print("✅ 组件初始化成功，开始集成性能测试...")
        
        monitor.start_monitoring()
        
        # 测试50次完整流程
        test_cycles = 50
        successful_cycles = 0
        
        for i in range(test_cycles):
            start_time = time.time()
            
            # 获取图像
            if use_camera:
                ret, frame = camera.cap.read()
                if not ret:
                    # 如果摄像头失败，使用模拟图像
                    frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
            else:
                # 使用模拟图像
                frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
                cv2.rectangle(frame, (100+i, 100+i), (200+i, 200+i), (255, 255, 255), -1)
            
            # YOLO检测
            detection_start = time.time()
            detections = detector.detect(frame)
            detection_time = time.time() - detection_start
            monitor.record_detection_time(detection_time)
            
            # 绘制检测结果
            processed_frame = detector.draw_detections(frame, detections)
            
            # 添加瞄准线
            final_frame = crosshair.draw(processed_frame)
            
            successful_cycles += 1
            
            cycle_time = time.time() - start_time
            monitor.record_frame_time(cycle_time)
            
            if (i + 1) % 10 == 0:
                print(f"进度: {i + 1}/{test_cycles} 个周期")
        
        monitor.stop_monitoring()
        
        if use_camera:
            camera.release()
        
        # 获取统计信息
        stats = monitor.get_statistics()
        
        print(f"\n📊 集成性能统计:")
        print(f"成功周期: {successful_cycles}/{test_cycles}")
        print(f"成功率: {successful_cycles/test_cycles*100:.1f}%")
        
        if 'frame_rate' in stats:
            fr = stats['frame_rate']
            print(f"整体帧率: {fr['avg_fps']:.1f} FPS")
            print(f"平均处理时间: {fr['avg_frame_time']:.1f} ms")
        
        if 'detection' in stats:
            det = stats['detection']
            print(f"检测时间占比: {det['avg_time']/fr['avg_frame_time']*100:.1f}%")
        
        if 'cpu' in stats:
            print(f"CPU使用率: 平均 {stats['cpu']['avg_usage']:.1f}%, 峰值 {stats['cpu']['max_usage']:.1f}%")
        
        if 'memory' in stats:
            print(f"内存使用率: 平均 {stats['memory']['avg_usage']:.1f}%, 峰值 {stats['memory']['max_usage']:.1f}%")
        
        return successful_cycles > test_cycles * 0.9  # 90%成功率
        
    except Exception as e:
        print(f"❌ 集成性能测试失败: {e}")
        return False

def main():
    """主测试函数"""
    print("🚀 开始性能测试")
    print("测试时间:", time.strftime("%Y-%m-%d %H:%M:%S"))
    
    results = []
    
    # 执行性能测试
    results.append(("摄像头性能", test_camera_performance()))
    results.append(("YOLO性能", test_yolo_performance()))
    results.append(("集成性能", test_integrated_performance()))
    
    # 输出测试结果
    print("\n" + "=" * 50)
    print("📊 性能测试结果汇总")
    print("=" * 50)
    
    passed = 0
    total = len(results)
    
    for test_name, result in results:
        status = "✅ 通过" if result else "❌ 失败"
        print(f"{test_name:12} : {status}")
        if result:
            passed += 1
    
    print("-" * 50)
    print(f"总计: {passed}/{total} 项性能测试通过")
    
    if passed == total:
        print("🎉 性能测试全部通过！系统性能良好")
        return True
    else:
        print("⚠️  部分性能测试未达标，建议优化")
        return False

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)